#include    <proc.h>
#include    <assert.h>

extern void _restart(struct C *prev, struct C *next) __attribute__ ((regparm(2)));
extern void *_get_eip(void);
extern void wrap_gate(void);

#define _get_esp \
({ \
    void *__esp__ = NULL; \
    asm volatile ("mov %%esp, %0" : "=r"(__esp__)); \
    __esp__; \
})

#define _get_ebp \
({ \
    void *__ebp__ = NULL; \
    asm volatile ("mov %%esp, %0" : "=r"(__ebp__)); \
    __ebp__; \
})


// context


static
struct C *make_C(unsigned stack_size) {
    struct C *cntx = 
        (struct C *) malloc(sizeof(struct C) + sizeof(void *) * stack_size);

    if (unlikely(cntx == NULL))
        return NULL;
    memset(cntx, 0, sizeof(struct C) + sizeof(void *) * stack_size);
    cntx->stack_size = stack_size;

    return cntx;
}

static
struct C *init_C(struct C *cntx, void (*entry)(void), void (*ret)(void)) {
    // XXX: should not crash
    assert(cntx != NULL);

    unsigned bottom = cntx->stack_size;

    // XXX: no ebp on frame
    // XXX: no preserved registers
    cntx->stack[bottom - 1] = ret;
    cntx->stack[bottom - 2] = entry;
    cntx->eip = wrap_gate;
    cntx->ebp = cntx->stack + bottom - 1;
    cntx->esp = cntx->stack + bottom - 2;

    return cntx;
}

static
void free_C(struct C *cntx) {
    assert(cntx != NULL);
    free(cntx);
}

#define _save(_curr_) \
    do { \
        _curr_->esp = _get_esp; \
        _curr_->ebp = _get_ebp; \
        _curr_->eip = _get_eip(); \
    } while (0)


// limbo


static
struct L *make_L(void (*entry)(void), unsigned stack_size) {
    struct L *limbo = (struct L *) malloc(sizeof(struct L));
    assert(limbo != NULL);

    limbo->entry = entry;
    limbo->cntx = init_C(make_C(stack_size), entry, NULL);
    assert(limbo->cntx != NULL);

    return limbo;
}

static
void free_L(struct L *limbo) {
    assert(limbo != NULL);
    free(limbo->cntx);
    free(limbo);
}


// cothread


static
struct T *make_T(void (*entry)(void), 
                 void (*ret)(void), 
                 void *arg, 
                 unsigned stack_size) {
    struct T *th = (struct T *) malloc(sizeof(struct T));
    assert(th != NULL);

    th->cntx = init_C(make_C(stack_size), entry, ret);
    assert(th->cntx != NULL);

    th->arg = arg;

    return th;
}

static
void free_T(struct T *th) {
    assert(th != NULL);
    free(th->cntx);
    free(th);
}


// carrier


static
struct M *make_M(unsigned index,
                 void (*limbo_entry)(void),
                 unsigned cntx_stack_size,
                 unsigned limbo_stack_size) {
    struct M *mth = (struct M *) malloc(sizeof(struct M));
    assert(mth != NULL);

    mth->cntx = make_C(cntx_stack_size);
    mth->limbo = make_L(limbo_entry, limbo_stack_size);
    assert(mth->cntx != NULL);
    assert(mth->limbo != NULL);

    mth->index = index;
    mth->running = NULL;
    pthread_mutex_init(&(mth->queue_mutex), NULL);
    init_list_head(&(mth->queue));

    return mth;
}


// runtime implementation


static __thread struct M *self_carrier = NULL;
static struct M *main_carrier = NULL;
static struct M **carrier_list = NULL;
static unsigned n_carriers = 1;
static pthread_mutex_t sched_mutex, sync_mutex;


static
void limbo_loop(void) {
    assert(self_carrier != NULL);

    while (1) {
        struct T *th = running_T(self_carrier);

        if (th != NULL) {
            pthread_mutex_lock(&(self_carrier->queue_mutex));
            list_del(&(th->lctl));
            self_carrier->running = NULL;
            pthread_mutex_unlock(&(self_carrier->queue_mutex));
            free_T(th);
        }
        if (self_carrier != main_carrier) {
            _restart(get_cntx(get_limbo(self_carrier)), get_cntx(self_carrier));
        } else {
            list_ctl_t *r = self_carrier->queue.next;
            struct T *run = 
                self_carrier->running = container_of(r, struct T, lctl);
            _restart(get_cntx(get_limbo(self_carrier)), get_cntx(run));
        }
    }
}

static
void exit_gate(void) {
    assert(self_carrier != NULL);
    assert(running_T(self_carrier) != NULL);
    _restart(get_cntx(running_T(self_carrier)), get_cntx(get_limbo(self_carrier)));
}

static
void *sched_loop(void *arg) {
    assert(arg != NULL);
    self_carrier = (struct M *) arg;

    // XXX: Sync barrier?
    pthread_mutex_lock(&sync_mutex);
    pthread_mutex_unlock(&sync_mutex);

    pthread_mutex_lock(&(self_carrier->queue_mutex));
    list_ctl_t *h = &(self_carrier->queue);
    pthread_mutex_unlock(&(self_carrier->queue_mutex));

    _save(get_cntx(self_carrier));

    list_ctl_t *p = h;
    list_ctl_t *bp = h->next;

    while (1) {
        pthread_mutex_lock(&(self_carrier->queue_mutex));
        p = bp;
        if (p == h)
            (p = p->next), (bp = bp->next);
        bp = bp->next;
        pthread_mutex_unlock(&(self_carrier->queue_mutex));

        //_save(get_cntx(self_carrier));

        if (p == h) {
            self_carrier->running = NULL;
            continue;
        } else {
            self_carrier->running = container_of(p, struct T, lctl);
        }

        _restart(get_cntx(self_carrier), get_cntx(self_carrier->running));
    }
}

static
struct T *sched_next(void) {
    static list_ctl_t *p = NULL, *bp = NULL;

    if (unlikely(p == NULL)) {
        p = &(main_carrier->queue);
        bp = p->next;
    }

    pthread_mutex_lock(&(main_carrier->queue_mutex));
    p = bp;
    if (p == &(main_carrier->queue))
        (p = p->next), (bp = bp->next);
    bp = bp->next;
    pthread_mutex_unlock(&(main_carrier->queue_mutex));

    struct T *next_th = NULL;

    if (p != &(main_carrier->queue)) {
        next_th = container_of(p, struct T, lctl);
    } else {
        next_th = container_of(main_carrier->queue.next, struct T, lctl);
    }

    return next_th;
}

static inline
unsigned get_next_carrier(int forbid) {
    static next_carrier = 0;

    pthread_mutex_lock(&(sched_mutex));
    next_carrier = (next_carrier + 1) % n_carriers;
    (forbid) && (next_carrier == 0) 
             && (next_carrier = (next_carrier + 1) % n_carriers);
    pthread_mutex_unlock(&(sched_mutex));

    return next_carrier;
}

static
void init_main_carrier(unsigned cntx_stack_size, unsigned limbo_stack_size) {
    main_carrier = make_M(0, limbo_loop, cntx_stack_size, limbo_stack_size);
    main_carrier->tid = pthread_self();
    //main_carrier->running = NULL;
    carrier_list[0] = main_carrier;
    main_carrier->running = (struct T *) malloc(sizeof(struct T));
#define SMALL_STACK 64
    main_carrier->running->cntx = 
        (struct C *) malloc(sizeof(struct C) + sizeof(void *) * SMALL_STACK);
    memset(main_carrier->running->cntx, 
           0, 
           sizeof(struct C) + sizeof(void *) * SMALL_STACK);
    main_carrier->running->arg = (void *) 0xdead;
    list_add_tail(&(main_carrier->running->lctl), &(main_carrier->queue));
#undef  SMALL_STACK
}

int viphyr_go(void (*func)(void), void *arg, unsigned stack_size, int forbid) {
    struct T *th = make_T(func, exit_gate, arg, stack_size);
    if (th == NULL)
        return -1;

    unsigned index = get_next_carrier(forbid);
    pthread_mutex_lock(&(carrier_list[index]->queue_mutex));
    list_add_tail(&(th->lctl), &(carrier_list[index]->queue));
    pthread_mutex_unlock(&(carrier_list[index]->queue_mutex));

    return 0;
}

void viphyr_resched(void) {
    if (self_carrier != main_carrier) {
        _restart(get_cntx(running_T(self_carrier)), get_cntx(self_carrier));
    } else {
        struct T *next = sched_next();
        if (next != NULL) {
            struct T *curr = running_T(self_carrier);
            self_carrier->running = next;
            if (curr != next) {
                assert(next->cntx != NULL);
                _restart(curr->cntx, next->cntx);
            }
        }
    }
}

int viphyr_init(unsigned more_carriers) {
    pthread_mutex_init(&sched_mutex, NULL);
    pthread_mutex_init(&sync_mutex, NULL);

    n_carriers += more_carriers;
    carrier_list = (struct M **) malloc(sizeof(struct M *) * n_carriers);
#define OMIT_SIZE   64
#define LIMBO_SIZE  256
    init_main_carrier(OMIT_SIZE, LIMBO_SIZE);
    self_carrier = main_carrier;

    if (more_carriers > 0) {
        pthread_mutex_lock(&sync_mutex);

        pthread_attr_t attr;
        pthread_attr_init(&attr);
        pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);

        int itr;
        for (itr = 1; itr < n_carriers; itr++) {
            carrier_list[itr] = make_M(itr, limbo_loop, OMIT_SIZE, LIMBO_SIZE);
            pthread_create(&(carrier_list[itr]->tid),
                           &attr,
                           sched_loop,
                           carrier_list[itr]);
        }

        pthread_mutex_unlock(&sync_mutex);
    }
#undef  OMIT_SIZE
#undef  LIMBO_SIZE

    return 0;
}


// debug


#include    <stdio.h>


void some_func(void) {
    printf("hello\n");
}

int main(int argc, char *argv[]) {
    unsigned n = atoi(argv[1]);
    viphyr_init(n);

    int i;

    for (i = 0; i < 3; i++) {
        viphyr_go(some_func, NULL, 4096, n ? 1 : 0);
        viphyr_resched();
    }
    printf("world\n");

    getchar();

    return 0;
}
